import { exec } from 'child_process';
import { promisify } from 'util';
const execAsync = promisify(exec);
export interface DocumentRecord {
id: number;
uri: string;
title: string;
mtime: number;
hash: string;
metadata?: string;
created_at: string;
}
export interface ChunkRecord {
id: number;
doc_id: number;
text: string;
section: string;
offset: number;
lang: string;
hash: string;
created_at: string;
}
export interface SearchResult {
text: string;
section: string;
uri: string;
title: string;
score: number;
}
export class SQLiteClient {
private containerName: string;
private dbPath: string;
constructor() {
this.containerName = 'mcp-sqlite-1';
this.dbPath = '/data/rag.db';
}
async query<T>(sql: string, params: any[] = []): Promise<T[]> {
try {
// Заменяем ? на параметры напрямую
let preparedSql = sql;
for (let i = 0; i < params.length; i++) {
const param = params[i];
const escapedParam = typeof param === 'string' ? `'${param.replace(/'/g, "''")}'` : String(param);
preparedSql = preparedSql.replace('?', escapedParam);
}
const command = `docker exec ${this.containerName} sqlite3 ${this.dbPath} "${preparedSql}"`;
const { stdout, stderr } = await execAsync(command);
if (stderr && !stderr.includes('Warning')) {
throw new Error(`SQLite error: ${stderr}`);
}
// Парсим результат
return this.parseSQLiteOutput(stdout);
} catch (error) {
console.error('Ошибка SQLite запроса:', error);
throw new Error(`Ошибка SQLite: ${error}`);
}
}
async execute(sql: string, params: any[] = []): Promise<void> {
try {
// Заменяем ? на параметры напрямую
let preparedSql = sql;
for (let i = 0; i < params.length; i++) {
const param = params[i];
const escapedParam = typeof param === 'string' ? `'${param.replace(/'/g, "''")}'` : String(param);
preparedSql = preparedSql.replace('?', escapedParam);
}
const command = `docker exec ${this.containerName} sqlite3 ${this.dbPath} "${preparedSql}"`;
const { stderr } = await execAsync(command);
if (stderr && !stderr.includes('Warning')) {
throw new Error(`SQLite error: ${stderr}`);
}
} catch (error) {
console.error('Ошибка SQLite выполнения:', error);
throw new Error(`Ошибка SQLite: ${error}`);
}
}
async getDocuments(filters?: Record<string, unknown>): Promise<DocumentRecord[]> {
let sql = 'SELECT * FROM docs';
const params: any[] = [];
if (filters && Object.keys(filters).length > 0) {
const conditions: string[] = [];
Object.entries(filters).forEach(([key, value]) => {
conditions.push(`JSON_EXTRACT(metadata, '$.${key}') = ?`);
params.push(value);
});
sql += ` WHERE ${conditions.join(' AND ')}`;
}
sql += ' ORDER BY created_at DESC';
return this.query<DocumentRecord>(sql, params);
}
async getDocument(uri: string): Promise<DocumentRecord | null> {
const sql = 'SELECT * FROM docs WHERE uri = ? LIMIT 1';
const results = await this.query<DocumentRecord>(sql, [uri]);
return results.length > 0 ? results[0] : null;
}
async addDocument(uri: string, title: string, content: string, metadata?: string): Promise<number> {
const mtime = Math.floor(Date.now() / 1000);
const hash = this.generateHash(content);
// Добавляем документ и получаем ID в одной транзакции
const insertDocSql = `
INSERT INTO docs (uri, title, mtime, hash, metadata, created_at)
VALUES (?, ?, ?, ?, ?, datetime('now'));
SELECT last_insert_rowid();
`;
const result = await this.query<{ 'last_insert_rowid()': number }>(insertDocSql, [uri, title, mtime, hash, metadata || null]);
if (!result || result.length === 0) {
throw new Error('Не удалось получить ID добавленного документа');
}
const docId = result[0]['last_insert_rowid()'];
// Разбиваем на чанки и добавляем их
const chunks = this.chunkText(content);
for (let i = 0; i < chunks.length; i++) {
const chunk = chunks[i];
const chunkHash = this.generateHash(chunk);
const insertChunkSql = `
INSERT INTO chunks (doc_id, text, section, offset, lang, hash, created_at)
VALUES (?, ?, ?, ?, ?, ?, datetime('now'))
`;
await this.execute(insertChunkSql, [
docId,
chunk,
'main',
i * 1000,
'auto',
chunkHash
]);
}
return docId;
}
async search(query: string, limit: number = 5): Promise<SearchResult[]> {
try {
// Используем FTS5 для полнотекстового поиска
const sql = `
SELECT
ch.text,
ch.section,
d.uri,
d.title,
1.0 as score
FROM chunks_fts c
JOIN chunks ch ON c.rowid = ch.id
JOIN docs d ON ch.doc_id = d.id
WHERE chunks_fts MATCH ?
ORDER BY ch.created_at DESC
LIMIT ?
`;
const results = await this.query<any>(sql, [query, limit]);
return results.map(r => ({
text: r.text,
section: r.section,
uri: r.uri,
title: r.title,
score: parseFloat(r.score) || 0
}));
} catch (error) {
console.error('Ошибка поиска:', error);
// Fallback к простому поиску по LIKE
return this.simpleSearch(query, limit);
}
}
private async simpleSearch(query: string, limit: number): Promise<SearchResult[]> {
const sql = `
SELECT
c.text,
c.section,
d.uri,
d.title,
1.0 as score
FROM chunks c
JOIN docs d ON c.doc_id = d.id
WHERE c.text LIKE ?
ORDER BY c.created_at DESC
LIMIT ?
`;
const results = await this.query<any>(sql, [`%${query}%`, limit]);
return results.map(r => ({
text: r.text,
section: r.section,
uri: r.uri,
title: r.title,
score: r.score
}));
}
async deleteDocument(uri: string): Promise<void> {
const doc = await this.getDocument(uri);
if (!doc) {
throw new Error('Документ не найден');
}
// Удаляем чанки (каскадное удаление настроено в БД)
await this.execute('DELETE FROM docs WHERE uri = ?', [uri]);
}
async getDocumentChunks(docId: number): Promise<ChunkRecord[]> {
const sql = 'SELECT * FROM chunks WHERE doc_id = ? ORDER BY offset';
return this.query<ChunkRecord>(sql, [docId]);
}
async updateDocument(uri: string, title: string, content: string, metadata?: string): Promise<void> {
const doc = await this.getDocument(uri);
if (!doc) {
throw new Error('Документ не найден');
}
const mtime = Math.floor(Date.now() / 1000);
const hash = this.generateHash(content);
// Обновляем документ
await this.execute(
'UPDATE docs SET title = ?, mtime = ?, hash = ?, metadata = ? WHERE uri = ?',
[title, mtime, hash, metadata || null, uri]
);
// Удаляем старые чанки
await this.execute('DELETE FROM chunks WHERE doc_id = ?', [doc.id]);
// Добавляем новые чанки
const chunks = this.chunkText(content);
for (let i = 0; i < chunks.length; i++) {
const chunk = chunks[i];
const chunkHash = this.generateHash(chunk);
await this.execute(`
INSERT INTO chunks (doc_id, text, section, offset, lang, hash, created_at)
VALUES (?, ?, ?, ?, ?, ?, datetime('now'))
`, [doc.id, chunk, 'main', i * 1000, 'auto', chunkHash]);
}
}
private chunkText(text: string, chunkSize: number = 1000): string[] {
const chunks: string[] = [];
// Разбиваем по строкам и параграфам
const lines = text.split('\n').filter(line => line.trim().length > 0);
let currentChunk = '';
for (const line of lines) {
const trimmedLine = line.trim();
// Если строка начинается с #, это заголовок - начинаем новый чанк
if (trimmedLine.startsWith('#')) {
if (currentChunk.trim()) {
chunks.push(currentChunk.trim());
}
currentChunk = trimmedLine + '\n';
continue;
}
// Если добавление строки превысит размер чанка, сохраняем текущий
if ((currentChunk + trimmedLine + '\n').length > chunkSize && currentChunk.trim().length > 0) {
chunks.push(currentChunk.trim());
currentChunk = trimmedLine + '\n';
} else {
currentChunk += trimmedLine + '\n';
}
}
// Добавляем последний чанк
if (currentChunk.trim()) {
chunks.push(currentChunk.trim());
}
// Если чанки слишком большие, разбиваем дальше
const finalChunks: string[] = [];
for (const chunk of chunks) {
if (chunk.length <= chunkSize) {
finalChunks.push(chunk);
} else {
// Разбиваем большие чанки по предложениям
const sentences = chunk.split(/[.!?]+/).filter(s => s.trim().length > 0);
let currentSentence = '';
for (const sentence of sentences) {
if ((currentSentence + sentence).length > chunkSize && currentSentence.length > 0) {
finalChunks.push(currentSentence.trim());
currentSentence = sentence;
} else {
currentSentence += (currentSentence ? ' ' : '') + sentence;
}
}
if (currentSentence.trim()) {
finalChunks.push(currentSentence.trim());
}
}
}
return finalChunks.filter(chunk => chunk.trim().length > 0);
}
private generateHash(text: string): string {
// Простая хеш-функция для демонстрации
let hash = 0;
for (let i = 0; i < text.length; i++) {
const char = text.charCodeAt(i);
hash = ((hash << 5) - hash) + char;
hash = hash & hash; // Convert to 32bit integer
}
return Math.abs(hash).toString(36);
}
private parseSQLiteOutput<T>(output: string): T[] {
if (!output.trim()) return [];
const lines = output.trim().split('\n');
const results: T[] = [];
for (const line of lines) {
if (line.trim()) {
// Простой парсинг - предполагаем, что поля разделены |
const fields = line.split('|');
const obj: any = {};
// Для простоты, создаем объект с числовыми ключами
fields.forEach((field, index) => {
obj[index] = field.trim();
});
// Также добавляем доступ по имени поля для удобства
if (fields.length === 1) {
// Если это результат SELECT last_insert_rowid()
obj['last_insert_rowid()'] = parseInt(fields[0]) || 0;
}
// Для документов, добавляем доступ по имени поля
if (fields.length >= 6) {
obj['id'] = parseInt(fields[0]) || 0;
obj['uri'] = fields[1];
obj['title'] = fields[2];
obj['mtime'] = parseInt(fields[3]) || 0;
obj['hash'] = fields[4];
obj['metadata'] = fields[5] || null;
obj['created_at'] = fields[6] || '';
}
// Для результатов поиска, добавляем доступ по имени поля
if (fields.length >= 5) {
obj['text'] = fields[0];
obj['section'] = fields[1];
obj['uri'] = fields[2];
obj['title'] = fields[3];
obj['score'] = parseFloat(fields[4]) || 0;
}
results.push(obj as T);
}
}
return results;
}
async isConnected(): Promise<boolean> {
try {
await this.query('SELECT 1');
return true;
} catch {
return false;
}
}
}